from common_import import *
import pandas as pd
from common_import import *
import lightgbm as lgb
from lightgbm import LGBMClassifier
from common_import import *
import pandas as pd
import lightgbm as lgb


def get_reg_model():
    # 分离特征数据和目标值
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    era_activity = pd.read_csv("data/ER_activity_training.csv")

    # 移除不必要的列（例如化合物ID）
    X_train = molecular_descriptor[constants.feature_295]
    y_train = era_activity["pIC50"]

    # 使用LightGBM训练模型
    model = lgb.LGBMRegressor(
        objective="regression",
        n_estimators=100,
        random_state=40,
        learning_rate=0.05,
        max_depth=15,
        verbose=-1,
    )
    model.fit(X_train, y_train)

    # 返回训练好的模型
    return model


def get_cla_model():
    """
    数据预处理和模型训练，不计算精确度，返回训练好的模型。
    """
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    X_train = molecular_descriptor[constants.feature_295]
    ADMET = pd.read_csv("data/ADMET_training.csv")
    y_train = ADMET[["Caco-2", "CYP3A4", "hERG", "HOB", "MN"]]
    # 分割数据集为训练集和验证集
    # 训练LightGBM分类器
    lgbm_models = {}
    for target in ["Caco-2", "CYP3A4", "hERG", "HOB", "MN"]:
        lgbm = LGBMClassifier(
            n_estimators=100,
            learning_rate=0.1,
            random_state=42,
            force_col_wise=True,
            verbose=-1,  # 禁用所有输出信息
        )
        lgbm.fit(X_train, y_train[target])
        lgbm_models[target] = lgbm
    # 返回训练好的模型和最大特征重要性
    return lgbm_models


reg_model = get_reg_model()
cla_model = get_cla_model()


def predict_with_REG(model, X):
    """
    使用给定的模型对输入数据 X 进行预测。
    """
    y_pred = model.predict(X)
    return y_pred


def predict_with_CLA(lgbm_models, X):
    """
    使用训练好的模型对输入数据进行预测，返回 n 个元素的列表，
    每个元素根据条件 "Caco-2" + "CYP3A4" + 1 - "hERG" + "HOB" + 1 - "MN" >= 3
    判断是否为1或0。
    """
    # 存储每个目标的预测标签
    predictions = []

    # 对每个目标进行预测
    for target in ["Caco-2", "CYP3A4", "hERG", "HOB", "MN"]:
        model = lgbm_models[target]
        # 获取预测标签
        labels = model.predict(X)
        predictions.append(labels)

    # 将每个目标的预测标签值组合成 n 行 5 列的矩阵
    predictions_matrix = np.stack(predictions, axis=1)

    # 计算符合条件的行
    final_predictions = []
    for row in predictions_matrix:
        condition_value = (
            row[0]  # "Caco-2"
            + row[1]  # "CYP3A4"
            + (1 - row[2])  # 1 - "hERG"
            + row[3]  # "HOB"
            + (1 - row[4])  # 1 - "MN"
        )
        if condition_value >= 3:
            final_predictions.append(1)
        else:
            final_predictions.append(0)

    return final_predictions


def combined_prediction(X):

    # 调用回归模型预测
    reg_predictions = predict_with_REG(reg_model, X)

    # 调用分类模型预测
    cla_predictions = predict_with_CLA(cla_model, X)

    # 将两者相乘并返回
    combined_results = np.multiply(reg_predictions, cla_predictions)
    return combined_results.tolist()


if __name__ == "__main__":
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    selected_features = molecular_descriptor.loc[
        constants.indice_50[:50], constants.feature_295
    ]
    print(selected_features.shape)
    result = combined_prediction(selected_features)
    print(len(result))
